# build_graph顾名思义应该是建立论文中描述的global-based图的。
import pickle
import argparse

parser = argparse.ArgumentParser()
# parser.add_argument('--dataset', default='diginetica', help='diginetica/Tmall/Nowplaying')
# ========= 修改内容 ============
parser.add_argument('--dataset', default='diginetica', help='diginetica/Tmall/Nowplaying/sample')
# diginetica数据集所需GPU内存较大，，使用sample数据集进行本机测试
# ========= 修改内容 ============
parser.add_argument('--sample_num', type=int, default=12)
opt = parser.parse_args()

dataset = opt.dataset
sample_num = opt.sample_num

# 读取所有会话数据
seq = pickle.load(open('datasets/' + dataset + '/all_train_seq.txt', 'rb'))

if dataset == 'diginetica':
    num = 43098
elif dataset == "Tmall":
    num = 40728
elif dataset == "Nowplaying":
    num = 60417
# ========= 修改内容 ============
elif dataset == "sample":
    num = 310
# ========= 修改内容 ============
else:
    num = 3

# k阶邻居关系列表
relation = []
# 邻居列表
neighbor = [] * num

all_test = set()

# 初始化，元素为字典类型，且adj1的下标表示itemId，字典的keys为adj1[index]邻居集合，values为对应邻居出现的次数
adj1 = [dict() for _ in range(num)]
# 初始化，元素为列表类型
adj = [[] for _ in range(num)]

# 对于all_train_seq，对其进行遍历。
for i in range(len(seq)):
    data = seq[i]
    # k指的应该是k阶邻居，以1阶邻居为例，添加data[j]和data[j+1]之间的关系（k阶关系是双向的），以list的形式保存
    # 循环完后，relation保存的是所有会话中所有item之间的k阶邻居关系。
    # 论文里这里k值设定为2？
    for k in range(1, 4):
        for j in range(len(data) - k):
            # 以二元组的形式存储k阶邻居
            relation.append([data[j], data[j + k]])
            relation.append([data[j + k], data[j]])

# 遍历relation，adj1是由dict组成的列表，统计relation中所有邻居关系出现的次数，结果由adj1保存。
# 出现次数是指在所有会话中，k阶邻域内，两个节点作为邻居出现的次数
for tup in relation:
    # 如果tup[1]是tup[0]的邻居，将k阶邻居关系出现次数+1，否则，将该k阶邻居出现的次数置为1
    if tup[1] in adj1[tup[0]].keys():
        adj1[tup[0]][tup[1]] += 1
    else:
        adj1[tup[0]][tup[1]] = 1

weight = [[] for _ in range(num)]

# 遍历该数据集中的所有item，对adj1进行排序，记录所有邻居adj与各邻居出现的次数weight
for t in range(num):
    # 将项目t对应的k阶邻居按照出现次数进行排序，降序排列
    x = [v for v in sorted(adj1[t].items(), reverse=True, key=lambda x: x[1])]
    # 保存t对应的邻居，且根据其从高到低出现的次数排序。这里v是一个元组，v[0]是对应的邻居，v[1]是邻居对应的出现次数
    adj[t] = [v[0] for v in x]
    # adj中邻居对应的权重（出现次数）
    weight[t] = [v[1] for v in x]

# 只保留前sample_num个邻居及其对应出现的次数，这里选择出现次数最多的前sample_num个邻居
for i in range(num):
    adj[i] = adj[i][:sample_num]
    weight[i] = weight[i][:sample_num]

# 将adj与weight分别保存为adj.pkl和num.pkl
pickle.dump(adj, open('datasets/' + dataset + '/adj_' + str(sample_num) + '.pkl', 'wb'))
pickle.dump(weight, open('datasets/' + dataset + '/num_' + str(sample_num) + '.pkl', 'wb'))
